package dev.alm.aiserviceintegration.config;

import dev.langchain4j.model.chat.listener.ChatModelErrorContext;
import dev.langchain4j.model.chat.listener.ChatModelListener;
import dev.langchain4j.model.chat.listener.ChatModelRequestContext;
import dev.langchain4j.model.chat.listener.ChatModelResponseContext;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.Map;
import java.util.UUID;

@Configuration
@Slf4j
public class ChatModelListenerConfig {

    @Bean
    ChatModelListener chatModelListener() {
        return new ChatModelListener() {

            @Override
            public void onRequest(ChatModelRequestContext requestContext) {
                log.info("进入了ChatModelListener的onRequest方法");
                String uuid = UUID.randomUUID().toString();
                requestContext.attributes().put("traceId", uuid);
            }

            @Override
            public void onResponse(ChatModelResponseContext responseContext) {
                log.info("进入了ChatModelListener的onResponse方法");
                Map<Object, Object> attributes = responseContext.attributes();
                String traceId = (String) attributes.get("traceId");
                log.info("traceId: {}", traceId);
            }

            @Override
            public void onError(ChatModelErrorContext errorContext) {
                Map<Object, Object> attributes = errorContext.attributes();
                String traceId = (String) attributes.get("traceId");
                log.info("traceId: {}", traceId);
                log.info("进入了ChatModelListener的onError方法");
            }
        };
    }

}
